#include <mm/paging.h>
#include <mm/page_allocator.h>

#include <utils/memutils.h>

#include <cpuid.h>

#define PML4E_IDX(addr) (((addr) >> 39) & 0x1ff)
#define PDPTE_IDX(addr) (((addr) >> 30) & 0x1ff)
#define PDE_IDX(addr) (((addr) >> 21) & 0x1ff)
#define PTE_IDX(addr) (((addr) >> 12) & 0x1ff)

static uintptr_t canonical_address_mask = 0;
static uint8_t virtual_address_bits = 0;
static uint8_t physical_address_bits = 0;

uint64_t *kernel_pagetable = NULL;

void mm_paging_init() {
	uint32_t tmp, eax;

	__cpuid(0x80000000, eax, tmp, tmp, tmp);
	if(eax >= 0x80000008) {
		__cpuid(0x80000008, eax, tmp, tmp, tmp);
		physical_address_bits = eax & 0xff;
		virtual_address_bits = (eax >> 8) & 0xff;
	} else {
		physical_address_bits = 32;
		virtual_address_bits = 36;
	}

	canonical_address_mask = ~0ull >> (64 - virtual_address_bits);

	kernel_pagetable = (uint64_t *)mm_pallocator_alloc4k(1);
	memset(kernel_pagetable, 0, 4096);

	mm_paging_map_physical_mem(kernel_pagetable);

	__asm__ volatile("movq %0, %%cr3" ::"r"(kernel_pagetable));
}

int mm_paging_map_memory_4k(uint64_t *pagetable, uintptr_t vaddr, uintptr_t paddr, uint64_t flags) {
	if((vaddr & 0xfff) || (paddr & 0xfff) || !(flags & PAGE_P)) {
		return -EINVAL;
	}

	if(!pagetable) {
		pagetable = kernel_pagetable;
	}

	uint16_t pml4e = PML4E_IDX(vaddr);
	uint16_t pdpte = PDPTE_IDX(vaddr);
	uint16_t pde = PDE_IDX(vaddr);
	uint16_t pte = PTE_IDX(vaddr);

	uint64_t *pml4 = pagetable;
	uint64_t *pdpt = NULL;
	uint64_t *pd = NULL;
	uint64_t *pt = NULL;

	if((pml4[pml4e] & PAGE_P) == 0) {
		pdpt = (uint64_t *)mm_pallocator_alloc4k(1);
		if(!pdpt) {
			return -ENOMEM;
		}

		memset(pdpt, 0, 4096);
		pml4[pml4e] = (uintptr_t)pdpt | PAGE_P | PAGE_RW | PAGE_US;
	} else {
		pdpt = (uint64_t *)(pml4[pml4e] & PT_ADDRMASK);
	}

	if((pdpt[pdpte] & PAGE_P) == 0 || (pdpt[pdpte] & BIT(7))) {
		pd = (uint64_t *)mm_pallocator_alloc4k(1);
		if(!pd) {
			return -ENOMEM;
		}

		memset(pd, 0, 4096);
		pdpt[pdpte] = (uintptr_t)pd | PAGE_P | PAGE_RW | PAGE_US;
	} else {
		pd = (uint64_t *)(pdpt[pdpte] & PT_ADDRMASK);
	}

	if((pd[pde] & PAGE_P) == 0 || (pd[pde] & BIT(7))) {
		pt = (uint64_t *)mm_pallocator_alloc4k(1);
		if(!pt) {
			return -ENOMEM;
		}

		memset(pt, 0, 4096);
		pd[pde] = (uintptr_t)pt | PAGE_P | PAGE_RW | PAGE_US;
	} else {
		pt = (uint64_t *)(pd[pde] & PT_ADDRMASK);
	}

	pt[pte] = (paddr & PT_ADDRMASK) | flags;

	uint64_t *current_pagetable = NULL;
	__asm__ volatile("movq %%cr3, %0" ::"r"(current_pagetable));

	if(current_pagetable == pagetable) {
		__asm__ volatile("invlpg (%0)" ::"r"(vaddr));
	}

	return 0;
}

int mm_paging_map_memory(uint64_t *pagetable, uintptr_t vaddr, uintptr_t paddr, uint64_t count_4k, uint64_t flags) {
	for(uint64_t i = 0; i < count_4k; ++i) {
		int stat = mm_paging_map_memory_4k(pagetable, vaddr + (i << 12), paddr + (i << 12), flags);

		if(stat != 0) {
			return stat;
		}
	}

	return 0;
}

uintptr_t mm_paging_get_paddr(uint64_t *pagetable, uintptr_t vaddr) {
	if(!pagetable || (vaddr & 0xfff)) {
		return 0;
	}

	uint16_t pml4e = PML4E_IDX(vaddr);
	uint16_t pdpte = PDPTE_IDX(vaddr);
	uint16_t pde = PDE_IDX(vaddr);
	uint16_t pte = PTE_IDX(vaddr);

	uint64_t *pml4 = pagetable;
	uint64_t *pdpt = NULL;
	uint64_t *pd = NULL;
	uint64_t *pt = NULL;

	if((pml4[pml4e] & PAGE_P) == 0) {
		return 0;
	} else {
		pdpt = (uint64_t *)(pml4[pml4e] & PT_ADDRMASK);
	}

	if((pdpt[pdpte] & PAGE_P) == 0) {
		return 0;
	} else {
		pd = (uint64_t *)(pdpt[pdpte] & PT_ADDRMASK);
	}

	if((pd[pde] & PAGE_P) == 0) {
		return 0;
	} else {
		pt = (uint64_t *)(pd[pde] & PT_ADDRMASK);
	}

	return (pt[pte] & PT_ADDRMASK);
}

void mm_paging_map_physical_mem(uint64_t *pagetable) {
	mm_paging_map_memory(pagetable, 0, 0, mm_pallocator_get_total_pages(), PAGE_P | PAGE_RW);
}